package com.shoulder.core.utils;


import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.BiPredicate;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class TreeUtils {


    /**
     * @param list           合成树的list
     * @param parentIdGetter 获取父id
     * @param idGetter       获取id
     * @param rootCheck      是否含根节点
     * @param setSubChildren 设置子属性
     * @param <T>            实体对象的父节点id类型
     * @param <E>            实体对象
     * @return List<E>
     */
    public static <T, E> List<E> build(List<E> list
            , Function<E, T> parentIdGetter
            , Function<E, T> idGetter
            , Predicate<E> rootCheck
            , BiConsumer<E, List<E>> setSubChildren) {

        List<E> result = new ArrayList<>();

        Map<T, List<E>> parentMap = new HashMap<>();

        for (E e : list) {
            T parentId = parentIdGetter.apply(e);
            parentMap.computeIfAbsent(parentId, k -> new ArrayList<>()).add(e);
        }

        for (E e : list) {
            setSubChildren.accept(e, parentMap.getOrDefault(idGetter.apply(e), new ArrayList<>()));
            if (rootCheck.test(e)) {
                result.add(e);
            }
        }

        return result;
    }


    /**
     * @param list           合成树的list
     * @param parentIdGetter 获取父id
     * @param idGetter       获取id
     * @param parentCheck    是否是父节点        对比上一个方法 这里是父节点 上一个方法是根节点
     * @param setChildren    设置子属性
     * @param <T>            实体对象的父节点id类型
     * @param <E>            实体对象
     * @return List<E>
     */
    public static <T, E> List<E> buildTree(List<E> list
            , Function<E, T> parentIdGetter
            , Function<E, T> idGetter
            , BiPredicate<T, T> parentCheck
            , BiConsumer<E, List<E>> setChildren) {

        List<E> parents = fetchParents(list, parentCheck, parentIdGetter, idGetter);

        List<String> collected = parents.stream().map(e -> idGetter.apply(e).toString()).collect(Collectors.toList());

        Map<T, List<E>> parentMap = new HashMap<>();

        List<E> result = new ArrayList<>();

        for (E e : list) {
            T parentId = parentIdGetter.apply(e);
            parentMap.computeIfAbsent(parentId, k -> new ArrayList<>()).add(e);
        }

        for (E e : list) {
            setChildren.accept(e, parentMap.getOrDefault(idGetter.apply(e), new ArrayList<>()));

            if (collected.contains(idGetter.apply(e).toString())) {
                result.add(e);
            }
        }


        return result;
    }


    /**
     * @param list           合成树的list
     * @param parentCheck    是否是父节点
     * @param parentIdGetter 获取父id
     * @param idGetter       获取id
     * @param <E>            实体对象
     * @param <T>            实体对象父节点id类型
     * @return
     */
    private static <E, T> List<E> fetchParents(List<E> list, BiPredicate<T, T> parentCheck, Function<E, T> parentIdGetter, Function<E, T> idGetter) {

        List<E> dataList = new ArrayList<>();
        for (E e : list) {
            E orElse = list.stream().filter(it -> parentCheck.test(idGetter.apply(it), parentIdGetter.apply(e))).findFirst().orElse(null);
            if (orElse == null) {
                dataList.add(e);
            }
        }
        return dataList;
    }

    /**
     * 将树平铺
     *
     * @param list           树列表
     * @param childrenGetter 获取子节点
     */
    public static <E> List<E> flatten(List<E> list, Function<E, List<E>> childrenGetter) {
        List<E> result = new ArrayList<>();
        Deque<E> stack = new ArrayDeque<>(list);
        while (!stack.isEmpty()) {
            E current = stack.pop();
            result.add(current);
            List<E> children = childrenGetter.apply(current);
            if (children != null) {
                for (int i = children.size() - 1; i >= 0; i--) {
                    stack.push(children.get(i));
                }
            }
        }
        return result;
    }


    /**
     * 获取子节点
     *
     * @param e              元素
     * @param idGetter       id函数
     * @param parentIdGetter parentId 函数
     * @param list           数据列表
     * @param resultList     返回结果
     */
    public static <E, T> List<E> sonList(E e, Function<E, T> idGetter, Function<E, T> parentIdGetter, List<E> list, List<E> resultList) {
        List<E> children = list.stream().filter(each -> parentIdGetter.apply(each).equals(idGetter.apply(e)))
                .peek(each -> {
                    sonList(each, idGetter, parentIdGetter, list, resultList);
                }).toList();

        if (!CollectionUtils.isEmpty(children)) {
            resultList.addAll(children);
        }

        return resultList;
    }


    /**
     * 获取绝对路径
     *
     * @param e              实体对象
     * @param idGetter       id函数
     * @param parentIdGetter 父id函数
     * @param resultGetter   结果返回函数
     * @param list           数据
     * @param rootCheck      父节点判断
     */
    public static <E, T, R> List<R> absolute(E e, Function<E, T> idGetter, Function<E, T> parentIdGetter, Function<E, R> resultGetter
            , List<E> list, Predicate<E> rootCheck) {
        //转为map对象
        Map<T, E> entityMap = list.stream().collect(Collectors.toMap(idGetter, Function.identity()));
        List<R> result = new ArrayList<>();
        return absolutePath(idGetter.apply(e), entityMap, parentIdGetter, result, resultGetter, rootCheck);
    }

    private static <T, E, R> List<R> absolutePath(T id, Map<T, E> entityMap, Function<E, T> parentIdGetter, List<R> result, Function<E, R> resultGetter, Predicate<E> rootCheck) {
        E e = entityMap.get(id);
        if (e == null) {
            return Collections.emptyList();
        }
        result.add(0, resultGetter.apply(e));
        T parentId = parentIdGetter.apply(e);
        if (rootCheck.test(e)) {
            return result;
        } else {
            return absolutePath(parentId, entityMap, parentIdGetter, result, resultGetter, rootCheck);
        }
    }
}

